# /usr/bin/env python3
# -*- coding: utf-8 -*-



import argparse
import sys
from ldap3 import Server, Connection, ALL, MODIFY_REPLACE
from sqlalchemy import create_engine, text

# --- 1. 配置信息 (请根据您的环境核对) ---

# LDAP 服务器配置
LDAP_SERVER_URI = 'ldap://127.0.0.1:389'
LDAP_BIND_DN = "cn=admin,dc=ai,dc=ustc,dc=edu,dc=cn"
LDAP_BIND_PASSWORD = "Ldap@UstCAi"
LDAP_BASE_DN = "dc=ai,dc=ustc,dc=edu,dc=cn"
LDAP_USER_SEARCH_BASE = f"ou=People,{LDAP_BASE_DN}"
LDAP_GROUP_SEARCH_BASE = f"ou=Group,{LDAP_BASE_DN}"

# 数据库文件路径 (假设在 instance 目录下)
DATABASE_URI = 'sqlite:///instance/app.db'

def get_approved_students_from_db(engine):
    """
    从数据库获取所有已批准学生及其导师的权威映射。
    使用正确的字符串条件进行查询。
    """
    print("Step 1: 从数据库加载已批准学生与导师的正确关系...")
    # 关键修正：使用字符串 'APPROVED' 和 'STUDENT' 进行查询
    query = """
        SELECT username, supervisor_name
        FROM application
        WHERE status = 'APPROVED' AND user_type = 'STUDENT';
    """
    try:
        with engine.connect() as connection:
            result = connection.execute(text(query))
            mapping = {row[0]: row[1] for row in result}
            print(f" -> 成功加载 {len(mapping)} 条已批准的学生记录。")
            return mapping
    except Exception as e:
        print(f"[错误] 无法连接或查询数据库: {e}")
        sys.exit(1)

def get_ldap_data(conn):
    """从LDAP获取所有用户和组的信息。"""
    print("Step 2: 从LDAP加载用户和组的当前状态...")
    
    # 获取组信息
    conn.search(LDAP_GROUP_SEARCH_BASE, '(objectClass=posixGroup)', attributes=['cn', 'gidNumber'])
    groups = conn.entries
    group_name_to_gid = {g.cn.value: int(g.gidNumber.value) for g in groups}
    gid_to_group_name = {int(g.gidNumber.value): g.cn.value for g in groups}

    # 获取用户信息
    conn.search(LDAP_USER_SEARCH_BASE, '(objectClass=posixAccount)', attributes=['uid', 'gidNumber'])
    users = {u.uid.value: {'gidNumber': int(u.gidNumber.value), 'dn': u.entry_dn} for u in conn.entries}

    print(f" -> 成功加载 {len(users)} 个用户和 {len(groups)} 个组。")
    return users, group_name_to_gid, gid_to_group_name

def main(dry_run):
    """主执行函数"""
    if dry_run:
        print("\n*** 正在以【演习模式】运行，不会对LDAP进行任何实际修改。 ***\n")
    else:
        print("\n*** 警告：正在以【执行模式】运行，将直接修改LDAP服务器数据！ ***\n")

    db_engine = create_engine(DATABASE_URI)
    approved_students = get_approved_students_from_db(db_engine)

    if not approved_students:
        print("数据库中没有找到已批准的学生，无需进行校验。")
        return

    try:
        server = Server(LDAP_SERVER_URI, get_info=ALL)
        conn = Connection(server, user=LDAP_BIND_DN, password=LDAP_BIND_PASSWORD, auto_bind=True)
    except Exception as e:
        print(f"[错误] 无法连接到LDAP服务器: {e}")
        sys.exit(1)
        
    try:
        ldap_users, group_name_to_gid, gid_to_group_name = get_ldap_data(conn)
        
        print("\nStep 3: 开始校验每个已批准学生的GID...\n" + "="*50)
        
        mismatch_count = 0
        for username, supervisor_name in approved_students.items():
            print(f"\n--- 正在校验用户: {username} (导师: {supervisor_name}) ---")

            if username not in ldap_users:
                print(f"  [警告] 用户 '{username}' 在数据库中已批准，但在LDAP中不存在。跳过。")
                continue

            # 确定正确的组和GID
            correct_group_name = f"group_{supervisor_name}"
            if correct_group_name not in group_name_to_gid:
                print(f"  [错误] 对应的导师组 '{correct_group_name}' 在LDAP中不存在。请先确保导师账户已正确配置。跳过。")
                continue
            
            correct_gid = group_name_to_gid[correct_group_name]
            
            # 获取当前的GID
            current_gid = ldap_users[username]['gidNumber']
            user_dn = ldap_users[username]['dn']

            # 对比并修正
            if current_gid == correct_gid:
                print(f"  [OK] GID正确。当前: {current_gid} ({correct_group_name})")
            else:
                mismatch_count += 1
                current_group_name = gid_to_group_name.get(current_gid, f"不存在的组(GID={current_gid})")
                print(f"  [不匹配] GID错误！")
                print(f"    - 当前 GID: {current_gid} (指向 -> {current_group_name})")
                print(f"    - 正确 GID: {correct_gid} (应该指向 -> {correct_group_name})")
                
                if not dry_run:
                    print(f"    -> [正在修正] 正在将用户 '{username}' 的gidNumber更新为 {correct_gid} ...")
                    conn.modify(user_dn, {'gidNumber': [(MODIFY_REPLACE, [str(correct_gid)])]})
                    if conn.result['result'] == 0:
                        print("    => [成功] 用户主组ID已更新。")
                    else:
                        print(f"    => [失败] 更新失败: {conn.result}")

        print("\n" + "="*50 + f"\n校验完成。共发现 {mismatch_count} 处不匹配。")

    finally:
        conn.unbind()
        print("已断开LDAP连接。")

if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description="校验并修正LDAP中已批准学生的gidNumber，使其与数据库中的导师信息保持一致。",
        formatter_class=argparse.RawTextHelpFormatter
    )
    parser.add_argument(
        '--execute',
        action='store_true',
        help="""实际执行LDAP修改操作。
如果没有此参数，脚本将以安全的“演习模式”(Dry Run)运行，只报告不匹配项。"""
    )
    args = parser.parse_args()

    main(dry_run=not args.execute)




















# # cleanup_home_dirs.py
# import os
# import shutil
# import argparse
# import sys
# from ldap3 import Server, Connection, ALL, MODIFY_REPLACE

# # --- 1. 配置信息 (请根据您的环境核对) ---

# # LDAP 服务器配置
# LDAP_SERVER_URI = 'ldap://127.0.0.1:389'
# LDAP_BIND_DN = "cn=admin,dc=ai,dc=ustc,dc=edu,dc=cn"
# LDAP_BIND_PASSWORD = "Ldap@UstCAi"
# # 只搜索用户所在的OU
# LDAP_USER_SEARCH_BASE = "ou=People,dc=ai,dc=ustc,dc=edu,dc=cn"

# # 要扫描的基础家目录
# HOME_BASE_DIR = '/home'

# # !! 关键 !! 安全排除列表：这些目录永远不会被删除
# # 请根据您的实际情况添加其他需要保留的目录名
# EXCLUDE_DIRS = [
#     'lost+found',
#     'user',       # 这个'user'目录似乎是您项目的一部分，务必保留
#     'admin',
#     'skel',
#     # 如果有其他系统账户或共享目录，请在这里添加
#     # 'public_share',
# ]


# def get_ldap_usernames():
#     """连接LDAP并获取所有用户的uid列表。"""
#     print("Step 1: 正在连接LDAP并获取所有合法用户名...")
#     try:
#         server = Server(LDAP_SERVER_URI, get_info=ALL)
#         conn = Connection(server, user=LDAP_BIND_DN, password=LDAP_BIND_PASSWORD, auto_bind=True)
        
#         # 只请求uid属性，提高效率
#         conn.search(LDAP_USER_SEARCH_BASE, '(objectClass=posixAccount)', attributes=['uid'])
        
#         # 使用集合(set)以获得极快的查找速度
#         usernames = {entry.uid.value for entry in conn.entries}
        
#         print(f" -> 成功从LDAP获取 {len(usernames)} 个用户名。")
#         conn.unbind()
#         return usernames
#     except Exception as e:
#         print(f"[错误] 无法连接或查询LDAP服务器: {e}")
#         sys.exit(1)

# def get_home_subdirs(base_path):
#     """获取指定路径下的所有子目录。"""
#     print(f"Step 2: 正在扫描 {base_path} 目录下的所有子目录...")
#     try:
#         all_items = os.listdir(base_path)
#         subdirs = [item for item in all_items if os.path.isdir(os.path.join(base_path, item))]
#         print(f" -> 发现 {len(subdirs)} 个子目录。")
#         return set(subdirs) #同样使用集合
#     except FileNotFoundError:
#         print(f"[错误] 基础目录 '{base_path}' 不存在。")
#         sys.exit(1)
#     except PermissionError:
#         print(f"[错误] 权限不足，无法读取 '{base_path}' 目录。请尝试使用 sudo 运行此脚本。")
#         sys.exit(1)

# def main(dry_run):
#     """主执行函数"""
#     # 获取权威的LDAP用户名列表
#     ldap_users = get_ldap_usernames()
    
#     # 获取实际存在的家目录列表
#     home_dirs = get_home_subdirs(HOME_BASE_DIR)
    
#     # 计算出需要排除的目录
#     exclude_set = set(EXCLUDE_DIRS)
    
#     # 找出孤儿目录：存在于home_dirs，但不在ldap_users且不在exclude_set中
#     orphan_dirs_names = home_dirs - ldap_users - exclude_set
    
#     # 转换为完整路径
#     orphan_dirs_paths = [os.path.join(HOME_BASE_DIR, name) for name in sorted(list(orphan_dirs_names))]

#     print("\nStep 3: 分析完成。")
    
#     if not orphan_dirs_paths:
#         print(" -> [优秀] 没有发现任何孤儿家目录。系统很干净！")
#         return

#     print("="*50)
#     if dry_run:
#         print("   --- 演习模式 (DRY RUN) ---   ")
#         print("以下孤儿目录【将会】在执行模式下被删除：")
#         for path in orphan_dirs_paths:
#             print(f"  - {path}")
#         print("\n要实际删除这些目录，请使用 --execute 参数运行脚本。")
#     else:
#         print("   --- 执行模式 (EXECUTE MODE) ---   ")
#         print("⚠️ 警告：以下目录及其所有内容将被【永久删除】！")
#         for path in orphan_dirs_paths:
#             print(f"  - {path}")
        
#         try:
#             # 最后一道安全防线
#             confirm = input("\n您确定要继续吗？请输入 'y' 确认: ")
#             if confirm.lower() != 'y':
#                 print("操作已取消。")
#                 return

#             print("\n正在执行删除操作...")
#             for path in orphan_dirs_paths:
#                 try:
#                     shutil.rmtree(path)
#                     print(f"  [成功] 已删除 {path}")
#                 except Exception as e:
#                     print(f"  [失败] 删除 {path} 时出错: {e}")
#         except KeyboardInterrupt:
#             print("\n操作被用户中断。")

#     print("="*50)


# if __name__ == '__main__':
#     parser = argparse.ArgumentParser(
#         description="清理/home下不存在于LDAP中的用户的家目录。",
#         formatter_class=argparse.RawTextHelpFormatter
#     )
#     parser.add_argument(
#         '--execute',
#         action='store_true',
#         help="""实际执行删除操作。
# 如果没有此参数，脚本将以安全的“演习模式”(Dry Run)运行，只报告将要删除的目录。"""
#     )
#     args = parser.parse_args()

#     main(dry_run=not args.execute)